from PIL import Image
import shutil
import os

anno_box_path = r"/home/pxxy2417/CelebA/Anno/list_bbox_celeba.txt"
claaification_path = r"/home/pxxy2417/CelebA/Eval/list_eval_partition.txt"
base_dir = "/home/pxxy2417/CelebA/Img/data"
img_dir = "/home/pxxy2417/CelebA/Img/img_celeba"
dataset_dir_list = [base_dir + "/train/images", base_dir + "/train/labels",
                    base_dir + "/val/images", base_dir + "/val/labels",
                    base_dir + "/test/images", base_dir + "/test/labels"]

if os.path.exists(base_dir):
    shutil.rmtree(base_dir)
for dir in dataset_dir_list:
    os.makedirs(dir)

train_num = 50000
val_num = 5000
test_num = 50
box_file_title = 39

box_file = open(anno_box_path, "r")
box_file.seek(box_file_title)


def process(url, num):
    i = 0
    for line in box_file:
        imgname = line[0:6]
        img_strs = line.split()
        x1, y1, w, h = int(img_strs[1]), int(img_strs[2]), int(img_strs[3]), int(img_strs[4])
        x2, y2 = x1 + w, int((y1 + h) * 0.88)
        # 微调一下矩形 ，有点大

        img = Image.open(f"{img_dir}/{img_strs[0]}")
        img_w, img_h = img.size

        dw = 1. / (int(img_w))
        dh = 1. / (int(img_h))
        x = ((x1 + x2) / 2.0) * dw
        y = ((y1 + y2) / 2.0) * dh
        w = (x2 - x1) * dw
        h = (y2 - y1) * dh

        label_txt = open(f"{url}/labels/{imgname}.txt", "w")
        label_txt.write(f"0 {x} {y} {w} {h}\n")
        label_txt.flush()
        label_txt.close()

        shutil.copy(f"{img_dir}/{img_strs[0]}", f"{url}/images/{img_strs[0]}")

        i += 1
        if i >= num:
            break


process(base_dir + "/train", train_num)
process(base_dir + "/val", val_num)
process(base_dir + "/test", test_num)

box_file.close()
